package ip

import (
	"encoding/binary"
)

// SplitIPHeader 将Header切割成元素为16位大小的数组
// ihl 为header的长度
func SplitIPHeader(buf []byte, ihl int) []uint16 {
	arr := make([]uint16, 0, 10)
	var i int
	for i < ihl-1 {
		if i == hCheck {
			arr = append(arr, 0)
		} else {
			arr = append(arr, binary.BigEndian.Uint16(buf[i:]))
		}
		i += 2
	}
	if i < ihl {
		arr = append(arr, binary.BigEndian.Uint16([]byte{0, buf[i]}))
	}
	return arr
}

func CalIPCheckSum(buf []byte, count int) uint16 {
	return CalCheckSum(SplitIPHeader(buf, count))
}

func CalCheckSum(arr []uint16) uint16 {
	var sum uint32
	for _, i := range arr {
		sum += uint32(i)
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return uint16(^sum)
}

func VerifyCheckSum(buf []byte) bool {
	var sum uint32
	var i int
	for i < len(buf)-1 {
		sum += uint32(binary.BigEndian.Uint16(buf[i:]))
		i += 2
	}
	if i < len(buf) {
		sum += uint32(binary.BigEndian.Uint16([]byte{0, buf[i]}))
	}

	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return uint16(^sum) == 0
}
